"""
下载参考文献
"""
import os
import time
import traceback

import fitz
import pymupdf4llm
import requests
from bs4 import BeautifulSoup
import pathlib
from conf import Config as myconf
from loguru import logger


def download_file(url, local_filename):
  # 发送HTTP GET请求
  with requests.get(url, stream=True) as r:
    r.raise_for_status()
    # 打开本地文件进行写操作
    with open(local_filename, 'wb') as f:
      for chunk in r.iter_content(chunk_size=8192):
        f.write(chunk)
  return local_filename


class DownloadRefPdf:
  url0s = ['https://sci-hub.se', 'https://sci-hub.st', 'https://sci-hub.ru']

  @classmethod
  def get_local_file_by_doi(cls, doi):
    doi = doi.replace('/', '-')
    local_file = f'{myconf.dl_dir}/{doi}.pdf'
    return local_file

  @classmethod
  def download_ref_pdf(cls, doi):
    def gen_pdf_url(index_html):
      soup = BeautifulSoup(index_html, 'html.parser')
      btn_dl = soup.select('#buttons > button')
      if len(btn_dl) > 0:
        if 'onclick' in btn_dl[0].attrs:
          # 如： location.href='//moscow.sci-hub.ru/2207/a07000873cd7b3725fbd26dd752f383e/xu2013.pdf?download=true'
          pdf_url = btn_dl[0]['onclick'].split("href=")[1].replace("'", "").replace('"', '')
          if pdf_url.startswith('//'):
            pdf_url = 'https:' + pdf_url
          return pdf_url

    for url0 in cls.url0s:
      qry_url = url0 + '/' + doi
      logger.info(f'qry_url:  {qry_url}')
      try:
        response = requests.get(qry_url)
        if response.status_code == 200:
          index_html = response.content
          pdf_url = gen_pdf_url(index_html)
          if pdf_url:
            local_file = cls.get_local_file_by_doi(doi)
            logger.info(f'pdf_url:  {pdf_url}  , begin download to local_file:  {local_file}')
            local_file = download_file(pdf_url, local_file)
            if local_file:
              logger.info(f'download success:  {local_file}')
              return local_file
            else:
              logger.error(f'download fail:  {pdf_url}')
          else:
            logger.error(f'not found pdf download button in {qry_url}')
      except Exception as e:
        logger.error(traceback.format_exc())

  @classmethod
  def download_ref_pdf2_by_doi(cls, doi):
    to_local_file = cls.get_local_file_by_doi(doi)
    for url0 in cls.url0s:
      qry_url = url0 + '/' + doi
      logger.info(f'qry_url:  {qry_url}')
      local_file = cls.download_ref_pdf2(qry_url, to_local_file)
      if local_file:
        return local_file

  @classmethod
  def download_ref_pdf2(cls, url, to_local_file=None):
    from playwright.sync_api import sync_playwright

    def find_PDF_DL_Button(page):
      pdf_btns = page.locator("#buttons > button").all()
      return pdf_btns

    with sync_playwright() as p:
      browser = p.chromium.connect_over_cdp("http://localhost:9222")
      # browser = p.firefox.connect_over_cdp("http://localhost:9222")
      try:
        default_context = browser.contexts[0]
        page = default_context.pages[0]
        page.goto(url, timeout=60000)
        page.wait_for_load_state("load")
        time.sleep(2)
        logger.info(f'page.title={page.title}')
        dl_pdf_btns = find_PDF_DL_Button(page)
        if len(dl_pdf_btns) > 0:
          dl_pdf_btn = dl_pdf_btns[0]
          if dl_pdf_btn.is_enabled():
            with page.expect_download() as download_info:
              dl_pdf_btn.click()
              download = download_info.value
              path = download.path()  # 等待下载完成
              local_file = to_local_file if to_local_file else '/tmp/' + download.suggested_filename
              download.save_as(local_file)
              logger.info(f"save file to {local_file} , and close download page : ")
              return local_file
      finally:
        logger.info("close browser")
        browser.close()

  @classmethod
  def pdf_to_md(cls, pdf_file):
    if pdf_file.endswith('.pdf') and os.path.exists(pdf_file):
      md_text = pymupdf4llm.to_markdown(pdf_file)
      md_file = pdf_file.replace('.pdf', '.md')
      pathlib.Path(md_file).write_bytes(md_text.encode())
      return md_file
    else:
      logger.error(f'pdf_file not exists or not pdf file:  {pdf_file}')
      return None

  @classmethod
  def pdf_to_txt(cls, pdf_file, page_number: int = 1):
    if pdf_file.endswith('.pdf') and os.path.exists(pdf_file):
      # 打开PDF文件
      doc = fitz.open(pdf_file)
      if page_number < 0:
        page_number = doc.page_count + page_number + 1

      # 检查页面号是否有效
      if page_number < 1 or page_number > len(doc):
        return "页面号超出文档范围"

      page = doc[page_number - 1]
      text = page.get_text()
      doc.close()
      return text

  @classmethod
  def pdf_all_to_txt(cls, pdf_file):
    if pdf_file.endswith('.pdf') and os.path.exists(pdf_file):
      doc = fitz.open(pdf_file)
      txt_list = []
      for i in range(len(doc)):
        page = doc[i]
        txt_list.append( page.get_text())
      doc.close()
      return "\n".join(txt_list)


class DRPdf(DownloadRefPdf):
  pass


# test code
def test1():
  doi = '10.1001/jama.2013.168118'
  doi = '10.1056/nejmp058318'
  doi = '10.1056/NEJMoa0908292'
  doi = '10.2337/diacare.20.11.1664'
  pdf_file = DRPdf.download_ref_pdf(doi)
  md_file = DRPdf.pdf_to_md(pdf_file)
  logger.info(f'pdf_file:  {pdf_file}  , md_file:  {md_file}')


def test2():
  doi = '10.1001/jama.2013.168118'
  doi = '10.1001/jama.2013.168118'
  url = f'https://sci-hub.se/{doi}'

  # DRPdf.download_ref_pdf2(url)

  local_file = DRPdf.download_ref_pdf2_by_doi(doi)
  logger.info("local_file = " + local_file)


def test_pdf_to_md():
  pdf_file = r"E:\tmp\docu_dl_analyze\Haibo Wang NEJM-Trends+in+CKD+in+China.pdf"
  # pdf_file = r"D:\tmp\0228\00\2col2tables.pdf"
  pdf_file = r"D:\tmp\0228\00\0726\science-example1.pdf"
  pdf_file = r"D:\tmp\0228\00\0726\10.2337-diacare.20.11.1664.pdf"
  md_file = DRPdf.pdf_to_md(pdf_file)
  logger.info(f'md_file:  {md_file}')


def test_pdf_to_txt():
  pdf_file = r"E:\tmp\docu_dl_analyze\Haibo Wang NEJM-Trends+in+CKD+in+China.pdf"
  page_number = -1
  text = DRPdf.pdf_to_txt(pdf_file, page_number)
  logger.info(f'text:  {text}')

def test_pdf_all_to_txt():
  pdf_file = r"E:\tmp\docu_dl_analyze\Haibo Wang NEJM-Trends+in+CKD+in+China.pdf"
  text = DRPdf.pdf_all_to_txt(pdf_file)
  logger.info(f'text:\n{text}')


if __name__ == '__main__':
  # test2()
  # test_pdf_to_md()
  # test_pdf_to_txt()
  test_pdf_all_to_txt()

  r"""
  1. C:\tools\chrome-win64\chrome.exe --remote-debugging-port=9222
  2. "C:\Program Files\Mozilla Firefox\firefox.exe -start-debugger-server 9222"  --> 不可以连接，放弃
  """
